package com.smart.speaker.websocket;

import com.alibaba.dashscope.audio.asr.recognition.Recognition;
import com.alibaba.dashscope.audio.asr.recognition.RecognitionParam;
import com.alibaba.dashscope.audio.asr.recognition.RecognitionResult;
import com.alibaba.dashscope.common.ResultCallback;
import com.smart.speaker.config.ModelConfig;
import com.smart.speaker.util.OpusUtils;
import io.github.jaredmdobson.concentus.OpusDecoder;
import io.github.jaredmdobson.concentus.OpusException;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.impl.GenericObjectPool;

import java.nio.ByteBuffer;
import java.util.concurrent.CountDownLatch;

@Slf4j
public class RecognitionHandler {

    // 默认音频采样率和通道数
    private static final int DEFAULT_SAMPLE_RATE = 16000;
    private static final int DEFAULT_CHANNELS = 1;

    private final GenericObjectPool<Recognition> pool;
    private final SessionState sessionState;
    private final OpusDecoder opusDecoder;

    private volatile Recognition recognizer;
    private volatile CountDownLatch stopLatch;

    public RecognitionHandler(
            GenericObjectPool<Recognition> pool,
            SessionState sessionState
    ) {
        this.pool = pool;
        this.sessionState = sessionState;
        try {
            this.opusDecoder = new OpusDecoder(DEFAULT_SAMPLE_RATE, DEFAULT_CHANNELS); // 采样率16kHz，单声道
        } catch (OpusException e) {
            log.error("初始化OpusDecoder失败");
            throw new RuntimeException(e);
        }
    }

    public void initRecognizer() {
        if (recognizer != null) {
            return;
        }
        try {
            recognizer = pool.borrowObject();
            stopLatch = new CountDownLatch(1);
            RecognitionParam param = RecognitionParam.builder()
                    .model("paraformer-realtime-v1")
                    .format("pcm") // 'pcm'、'wav'、'opus'、'speex'、'aac'、'amr', you
                    .sampleRate(16000) // supported 8000、16000
                    .apiKey(ModelConfig.Key.apiKey)
                    .build();
            ResultCallback<RecognitionResult> callback = new ResultCallback<>() {
                @Override
                public void onEvent(RecognitionResult message) {
                    sessionState.updateLastSpeakTime();
                    if (message.isSentenceEnd()) {
                        sessionState.appendRequest(message.getSentence().getText());
                        log.info("[process] Fix:{}", message.getSentence().getText());
                    } else {
                        log.info("[process] Result: {}", message.getSentence().getText());
                    }
                }

                @Override
                public void onComplete() {
                    log.info("recognition complete");
                    pool.returnObject(recognizer);
                    recognizer = null;
                    stopLatch.countDown();
                }

                @Override
                public void onError(Exception e) {
                    log.error("recognition error", e);
                    close();
                    stopLatch.countDown();
                }
            };
            recognizer.call(param, callback);
        } catch (Exception e) {
            log.error("speaker init recognizer error", e);
            if (recognizer != null) {
                close();
            }
        }
    }

    public void sendAudioFrame(byte[] bytes) {
        try {
            byte[] pcmBytes = OpusUtils.opusFrameToPcm(opusDecoder, bytes);
            recognizer.sendAudioFrame(ByteBuffer.wrap(pcmBytes));
        } catch (Exception e) {
            log.error("speaker send audio frame error", e);
            close();
        }
    }

    public void stop() {
        if (recognizer != null) {
            try {
                recognizer.stop();
                stopLatch.await();
            } catch (Exception e) {
                log.error("speaker stop recognizer error", e);
                close();
            }
        }
    }

    private void close() {
        if (recognizer != null) {
            try {
                recognizer.getDuplexApi().close(1000, "bye");
                pool.invalidateObject(recognizer);
                pool.addObject();
            } catch (Exception e) {
                log.error("speaker close recognizer error", e);
            }
            recognizer = null;
        }
    }
}
